from segment_anything import sam_model_registry, SamPredictor
from segment_anything.samUtils import device, model_type, sam_checkpoint


class SingletonMeta(type):
    """
    这是一个单例元类，用于创建一个单例类。
    """

    _instances = {}

    def __call__(cls, *args, **kwargs):
        if cls not in cls._instances:
            instance = super().__call__(*args, **kwargs)
            cls._instances[cls] = instance
        return cls._instances[cls]


class PredictorPool(metaclass=SingletonMeta):
    def __init__(self, max_size=1):
        if not hasattr(self, 'initialized'):  # 防止初始化代码被多次执行
            self.predictors = []
            self.max_size = max_size
            self.initialized = True

    def get_predictor(self, image):
        # 如果有可用的predictor，直接返回
        for predictor in self.predictors:
            if not predictor.is_busy:
                predictor.set_image(image)
                predictor.is_busy = True
                return predictor

        # 如果没有可用的predictor且未达到最大数量，创建新的predictor
        if len(self.predictors) < self.max_size:
            new_predictor = self.create_predictor(image)
            self.predictors.append(new_predictor)
            return new_predictor

        # 如果达到最大数量，重用最早的predictor
        oldest_predictor = self.predictors.pop(0)
        oldest_predictor.set_image(image)
        oldest_predictor.is_busy = True
        self.predictors.append(oldest_predictor)
        return oldest_predictor

    def getPredictor(self):
        oldest_predictor = self.predictors.pop(1)
        oldest_predictor.is_busy = True
        self.predictors.append(oldest_predictor)
        return oldest_predictor

    def create_predictor(self, image):
        # 创建并返回一个新的predictor实例
        sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
        sam.to(device=device)
        predictor = SamPredictor(sam)
        predictor.set_image(image)
        predictor.is_busy = True  # 标记predictor正在使用
        return predictor

    def release_predictor(self, predictor):
        # 标记predictor为未使用状态
        predictor.is_busy = False
